from pathlib import Path
from openai import OpenAI
import json
 
client = OpenAI(
    api_key="sk-CwhcRSwMALExKTPwtKvsROiaHfmGh8jIQLeqYITzdweszJqB",
    base_url = "https://api.moonshot.cn/v1",
    # api_key= "sk-8c2efa282764429096ad74da241211a8",
    # base_url="https://dashscope.aliyuncs.com/compatible-mode/v1",
)
subfile = "./subtitle.srt"  # 修改文件路径为当前目录

# file_object = client.files.create(file=Path(subfile), purpose="file-extract")
# file_content = client.files.content(file_id=file_object.id).text

with open(subfile, "r", encoding='utf-8') as f:
    data = json.load(f)

# 读取prompt文件
with open("./prompt/subtitle.pt", "r", encoding='utf-8') as pf:
    prompt_content = pf.read()

completion = client.chat.completions.create(
    model="moonshot-v1-128k",
    max_tokens=10240,
    messages=[
        {"role": "system", "content": prompt_content},
        {
            "role": "user",
            "content": json.dumps(data, indent=4)
        }
    ],
    temperature=0,
)

print(completion.choices[0].message.content)
# 保存处理后的字幕
processed_data = completion.choices[0].message.content
with open("processed_subtitle.json", "w", encoding='utf-8') as f:
    json.dump(processed_data, f, ensure_ascii=False, indent=4)

# completion = client.chat.completions.create(
#     model="moonshot-v1-8k",
#     messages=[
#         {"role": "system",
#          "content": "你是 Kimi，由 Moonshot AI 提供的人工智能助手，你更擅长中文和英文的对话。你会为用户提供安全，有帮助，准确的回答。同时，你会拒绝一切涉及恐怖主义，种族歧视，黄色暴力等问题的回答。Moonshot AI 为专有名词，不可翻译成其他语言。"},
#         {"role": "user", "content": "你好，请给出python库lxml的介绍以及详细用法"}
#     ],
#     temperature=0.3,
# )

# print(completion.choices[0].message)
